package com.example.calcite.mysql;

import com.example.calcite.ITable;
import org.apache.calcite.DataContext;
import org.apache.calcite.linq4j.Enumerable;
import org.apache.calcite.linq4j.Linq4j;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.schema.impl.AbstractTable;
import org.apache.calcite.sql.type.SqlTypeName;

import java.sql.*;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;

public class MySQLTable extends AbstractTable implements ITable {
    private final String jdbcUrl;
    private final Properties properties;
    private final String tableName;
    private String[] columns;
    private RelDataType rowType;

    public MySQLTable(String jdbcUrl, Properties properties, String tableName) {
        this.jdbcUrl = jdbcUrl;
        this.properties = properties;
        this.tableName = tableName;
        loadMetadata();
    }

    private void loadMetadata() {
        try (Connection connection = DriverManager.getConnection(jdbcUrl, properties)) {
            DatabaseMetaData metaData = connection.getMetaData();
            
            // 获取列信息
            ResultSet columns = metaData.getColumns(null, null, tableName, null);
            List<String> columnNames = new ArrayList<>();
            List<SqlTypeName> columnTypes = new ArrayList<>();
            
            while (columns.next()) {
                columnNames.add(columns.getString("COLUMN_NAME"));
                int sqlType = columns.getInt("DATA_TYPE");
                columnTypes.add(getSqlTypeName(sqlType));
            }
            
            this.columns = columnNames.toArray(new String[0]);
            
            // 创建行类型
            RelDataTypeFactory typeFactory = new org.apache.calcite.jdbc.JavaTypeFactoryImpl();
            RelDataTypeFactory.Builder builder = typeFactory.builder();
            for (int i = 0; i < this.columns.length; i++) {
                builder.add(columnNames.get(i), typeFactory.createSqlType(columnTypes.get(i)));
            }
            this.rowType = builder.build();
            
        } catch (SQLException e) {
            throw new RuntimeException("Failed to load table metadata", e);
        }
    }

    private SqlTypeName getSqlTypeName(int sqlType) {
        switch (sqlType) {
            case Types.INTEGER:
            case Types.BIGINT:
                return SqlTypeName.INTEGER;
            case Types.DECIMAL:
            case Types.NUMERIC:
                return SqlTypeName.DECIMAL;
            case Types.DOUBLE:
            case Types.FLOAT:
                return SqlTypeName.DOUBLE;
            case Types.DATE:
                return SqlTypeName.DATE;
            case Types.TIMESTAMP:
                return SqlTypeName.TIMESTAMP;
            case Types.BOOLEAN:
                return SqlTypeName.BOOLEAN;
            default:
                return SqlTypeName.VARCHAR;
        }
    }

    @Override
    public String[] getColumns() {
        return columns;
    }

    @Override
    public void insert(Object... values) {
        throw new UnsupportedOperationException("Insert operation is not supported for MySQL tables");
    }

    @Override
    public RelDataType getRowType(RelDataTypeFactory typeFactory) {
        return rowType;
    }

   /* @Override
    public Enumerable<Object[]> scan(DataContext root) {
        try (Connection connection = DriverManager.getConnection(jdbcUrl, properties);
             Statement statement = connection.createStatement();
             ResultSet resultSet = statement.executeQuery("SELECT * FROM " + tableName)) {
            
            List<Object[]> rows = new ArrayList<>();
            while (resultSet.next()) {
                Object[] row = new Object[columns.length];
                for (int i = 0; i < columns.length; i++) {
                    row[i] = resultSet.getObject(i + 1);
                }
                rows.add(row);
            }
            return Linq4j.asEnumerable(rows);
            
        } catch (SQLException e) {
            throw new RuntimeException("Failed to scan table", e);
        }
    }*/

    @Override
    public Enumerable<Object[]> scan(DataContext root, List<RexNode> filters) {
        try (Connection connection = DriverManager.getConnection(jdbcUrl, properties);
             Statement statement = connection.createStatement();
             ResultSet resultSet = statement.executeQuery("SELECT * FROM " + tableName)) {

            List<Object[]> rows = new ArrayList<>();
            while (resultSet.next()) {
                Object[] row = new Object[columns.length];
                for (int i = 0; i < columns.length; i++) {
                    row[i] = resultSet.getObject(i + 1);
                }
                rows.add(row);
            }
            return Linq4j.asEnumerable(rows);

        } catch (SQLException e) {
            throw new RuntimeException("Failed to scan table", e);
        }
    }
}